import os
from sentence_transformers import SentenceTransformer, InputExample, losses, models
from torch.utils.data import DataLoader
from sentence_transformers.evaluation import TripletEvaluator
from sentence_transformers.util import cos_sim
data = {
    "九寨沟": [
        "九寨沟以其多彩的湖泊、瀑布和雪山而闻名。",
        "作为世界自然遗产，九寨沟的五花海和长海是必看景点。",
        "这个位于阿坝州的景区是徒步和摄影爱好者的天堂。",
    ],
    "乐山大佛": [
        "乐山大佛是一尊巨大的石刻弥勒佛坐像，濒临岷江。",
        "游客可以乘船远观大佛全景，也可以近距离攀爬参观。",
        "这尊佛像开凿于唐代，历时约90年才最终完成。",
    ],
    "成都大熊猫基地": [
        "成都大熊猫繁育研究基地是观赏中国国宝的最佳地点。",
        "在这里，你可以看到从幼崽到成年大熊猫的各种形态。",
        "该基地致力于大熊猫的保护和繁育工作。",
    ]
}
print("正在加载微调后的模型...")
fine_tuned_model = SentenceTransformer('./sichuan_attractions_model')
# 准备我们的知识库（这里用所有句子模拟）
knowledge_base_sentences = []
for sentences in data.values():
    knowledge_base_sentences.extend(sentences)

print("正在将知识库向量化...")
knowledge_base_embeddings = fine_tuned_model.encode(knowledge_base_sentences, convert_to_tensor=True,
                                                    show_progress_bar=True)
# 模拟用户查询
user_query = "我想看可爱的大熊猫宝宝"
print(f"\n用户查询: '{user_query}'")
# 将用户查询也向量化
query_embedding = fine_tuned_model.encode(user_query, convert_to_tensor=True)
# 使用余弦相似度计算查询向量与知识库中所有向量的相似度
similarities = cos_sim(query_embedding, knowledge_base_embeddings)
# 找出最相似的句子
top_results = similarities[0].topk(1)  # 找出最相似的1个结果
print("\n--- 推荐结果 (最相关的描述) ---")
for i, (score, idx) in enumerate(zip(top_results.values, top_results.indices)):
    print(f"Top {i + 1} (相似度: {score:.4f}):")
    print(f"  -> {knowledge_base_sentences[idx]}")